package cn.yx.common.core.util.security;

import cn.hutool.core.codec.Base64;
import cn.yx.common.core.constant.message.security.AesUtilMessageKey;
import cn.yx.common.core.enums.CharsetEnum;
import cn.yx.common.core.enums.security.AesAlgorithmEnum;
import cn.yx.common.core.exception.CommonException;
import lombok.extern.slf4j.Slf4j;
import org.bouncycastle.jce.provider.BouncyCastleProvider;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.UnsupportedEncodingException;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.Security;
import java.security.spec.AlgorithmParameterSpec;
import java.util.Objects;

/**
 * <p>AES加密解密类</p>
 *
 * @author Wgssmart
 */
@Slf4j
public final class AesUtil {

    private AesUtil() {
    }

    /**
     * AES加密
     *
     * @param keyBytes         key字节数组
     * @param data             需要加密的数据
     * @param aesAlgorithmEnum AES算法
     * @return 加密后的UTF8编码字符串
     */
    public static String encode(byte[] keyBytes, String data, AesAlgorithmEnum aesAlgorithmEnum) {
        Cipher cipher = getCipher(keyBytes, null, aesAlgorithmEnum, Cipher.ENCRYPT_MODE);
        try {
            return new String(cipher.doFinal(Base64.decode(data)), CharsetEnum.UTF8.getValue());
        } catch (IllegalBlockSizeException | BadPaddingException | UnsupportedEncodingException e) {
            log.error("aes encode error: {}, algorithm: {}", e.getMessage(), aesAlgorithmEnum.getValue());
            throw new CommonException(AesUtilMessageKey.AES_ENCODE_ERROR, e);
        }
    }

    /**
     * AES解密
     *
     * @param keyBytes         key字节数组
     * @param data             需要加密的数据
     * @param aesAlgorithmEnum AES算法
     * @return 解密后的UTF8编码字符串
     */
    public static String decode(byte[] keyBytes, String data, AesAlgorithmEnum aesAlgorithmEnum) {
        Cipher cipher = getCipher(keyBytes, null, aesAlgorithmEnum, Cipher.DECRYPT_MODE);
        try {
            return new String(cipher.doFinal(Base64.decode(data)), CharsetEnum.UTF8.getValue());
        } catch (IllegalBlockSizeException | BadPaddingException | UnsupportedEncodingException e) {
            log.error("aes decode error: {}, algorithm: {}", e.getMessage(), aesAlgorithmEnum.getValue());
            throw new CommonException(AesUtilMessageKey.AES_ENCODE_ERROR, e);
        }
    }

    /**
     * AES加密
     *
     * @param keyBytes         key字节数组
     * @param ivBytes          iv向量数组，可为null
     * @param data             需要加密的数据
     * @param aesAlgorithmEnum AES算法
     * @return 加密后的UTF8编码字符串
     */
    public static String encode(byte[] keyBytes, byte[] ivBytes, String data, AesAlgorithmEnum aesAlgorithmEnum) {
        Cipher cipher = getCipher(keyBytes, ivBytes, aesAlgorithmEnum, Cipher.ENCRYPT_MODE);
        try {
            return new String(cipher.doFinal(Base64.decode(data)), CharsetEnum.UTF8.getValue());
        } catch (IllegalBlockSizeException | BadPaddingException | UnsupportedEncodingException e) {
            log.error("aes encode error: {}, algorithm: {}", e.getMessage(), aesAlgorithmEnum.getValue());
            throw new CommonException(AesUtilMessageKey.AES_ENCODE_ERROR, e);
        }
    }

    /**
     * AES解密
     *
     * @param keyBytes         key字节数组
     * @param ivBytes          iv向量数组，可为null
     * @param data             需要加密的数据
     * @param aesAlgorithmEnum AES算法
     * @return 解密后的UTF8编码字符串
     */
    public static String decode(byte[] keyBytes, byte[] ivBytes, String data, AesAlgorithmEnum aesAlgorithmEnum) {
        Cipher cipher = getCipher(keyBytes, ivBytes, aesAlgorithmEnum, Cipher.DECRYPT_MODE);
        try {
            return new String(cipher.doFinal(Base64.decode(data)), CharsetEnum.UTF8.getValue());
        } catch (IllegalBlockSizeException | BadPaddingException | UnsupportedEncodingException e) {
            log.error("aes decode error: {}, algorithm: {}", e.getMessage(), aesAlgorithmEnum.getValue());
            throw new CommonException(AesUtilMessageKey.AES_DECODE_ERROR, e);
        }
    }

    /**
     * 获取Cipher对象
     *
     * @param keyBytes         key字节数组
     * @param ivBytes          iv向量字节数组，可为null
     * @param aesAlgorithmEnum AES算法模式
     * @param cipherMode       加密或解密模式，加密1，解密2，参考CipherMode类
     * @return Cipher对象
     */
    public static Cipher getCipher(byte[] keyBytes, byte[] ivBytes, AesAlgorithmEnum aesAlgorithmEnum, int cipherMode) {
        Security.addProvider(new BouncyCastleProvider());
        AlgorithmParameterSpec ivSpec = Objects.nonNull(ivBytes) ? new IvParameterSpec(ivBytes) : null;
        try {
            Cipher cipher = Cipher.getInstance(aesAlgorithmEnum.getValue());
            SecretKeySpec keySpec = new SecretKeySpec(keyBytes, AesAlgorithmEnum.AES.getValue());
            if (Objects.nonNull(ivSpec)) {
                cipher.init(cipherMode, keySpec, ivSpec);
            } else {
                cipher.init(cipherMode, keySpec);
            }
            return cipher;
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException |
                 InvalidAlgorithmParameterException e) {
            log.error("getCipher error: {}, algorithm: {}, cipher mode: {}", e.getMessage(), aesAlgorithmEnum.getValue(), cipherMode);
            throw new CommonException(AesUtilMessageKey.GET_CIPHER_ERROR, e);
        }
    }

}
